package com.gm.math;

import Jama.Matrix;

/**
 * Created by Mouse on 2017/2/10.
 */


public class KalmanFilter {

    int n;//状态维度
    int m;//测量维度
    int k;//输入控制维度

    /**
     * 状态转移矩阵 [n][n]
     */
    Matrix A;

    /**
     * 输入控制到状态转移矩阵 [n][k]
     */
    Matrix B;

    /**
     * 状态到观测转换矩阵 [m][n]
     */
    Matrix H;

    /**
     * 状态误差协方差矩阵 [n][n]
     */
    Matrix P;

    /**
     * 预测噪声协方差矩阵 [n][n]
     */
    Matrix Q;

    /**
     * 测量噪声协方差矩阵 [m][m]
     */
    Matrix R;

    /**
     * k时刻的状态 [n][1]
     */
    Matrix x_k;

    /**
     * K时刻外界对系统的作用 [k][1]
     */
    Matrix u_k;

    /**
     * K时刻的观测值 [m][1]
     */
    Matrix z_k;

    /**
     * K时刻的kalman增益 [m][n]
     */
    Matrix kg_k;


    public KalmanFilter(int n, int m, int k) {
        this.n = n;
        this.m = m;
        this.k = k;
    }


    /**
     * x(k) = A(k)*x(k-1) + B(k)*u(k) + w(k)           预估值
     * w(k) ~ N(0,Q(k))
     * <p>
     * z(k) = H(k)*x(k) + v(k)                         观测值
     * v(k) ~ N(0,R(k))
     * <p>
     * x(k|k-1) = A*x(k-1|k-1)+B*u(k)                状态预测
     * P(k|k-1) = A*P(k-1|k-1)*A'+Q                  均方误差预测
     * x(k|k) = x(k|k-1) + kg(k)*(z(k) - H*x(k|k-1)) 最优值估算
     * kg(k) = P(k|k-1)*H'/(H*P(k|k-1)*H' + R)       滤波增益
     * P(k|k) = (I-kg(k)*H)*P(k|k-1)                 均方误差更新
     */
    public void filter(Matrix z_k) {

        checkNull();

        Matrix predict_x;

        //  x(k|k-1) = A*x(k-1|k-1)+B*u(k)
        predict_x = A.times(x_k).plus(B.times(u_k));

        //  P(k|k-1) = A*P(k-1|k-1)*A'+Q
        P = A.times(P).times(A.transpose()).plus(Q);

        //  kg(k) = P(k|k-1)*H'/(H*P(k|k-1)*H' + R)
        kg_k = P.times(H.transpose()).times((H.times(P).times(H.transpose()).plus(R)).inverse());

        //  x(k|k) = x(k|k-1) + kg(k)*(z(k) - H*x(k|k-1))
        x_k = predict_x.plus(kg_k.times(z_k.minus(H.times(predict_x))));

        //  P(k|k) = (I-kg(k)*H)*P(k|k-1)
        P = P.minus(kg_k.times(H).times(P));

    }

    public void update(Matrix z_k) throws Exception {
        checkZ(z_k);
        filter(z_k);
    }

    private void checkZ(Matrix z_k) throws Exception {
        if (z_k == null || z_k.getRowDimension() != m)
            throw new Exception("输入错误");
    }

    private void checkNull() {

        if (A == null)
            A = new Matrix(n, n);
        if (B == null)
            B = new Matrix(n, k);
        if (H == null)
            H = new Matrix(m, n);
        if (P == null)
            P = new Matrix(n, n);

        if (Q == null)
            Q = new Matrix(n, n);
        if (R == null)
            R = new Matrix(m, m);

        if (x_k == null)
            x_k = new Matrix(n, 1);
        if (u_k == null)
            u_k = new Matrix(k, 1);
    }


    public void setA(Matrix a) {
        A = a;
    }

    public void setB(Matrix b) {
        B = b;
    }

    public void setH(Matrix h) {
        H = h;
    }

    public void setP(Matrix p) {
        P = p;
    }

    public void setU_k(Matrix u_k) {
        this.u_k = u_k;
    }

    public void setQ(Matrix q) {
        Q = q;
    }

    public void setR(Matrix r) {
        R = r;
    }

    public Matrix getKg_k() {
        return kg_k;
    }

    public void setX_k(Matrix x_k) {
        this.x_k = x_k;
    }

    public Matrix getX_k() {
        return x_k;
    }

    public void setZ_k(Matrix z_k) {
        this.z_k = z_k;
    }

    public Matrix getZ_k() {
        return z_k;
    }
}
